# Usage: 存储处理面部腺样体肥大的功能函数
# Author：zuoyuwei
# Date: 2023/02/24
import cv2
import math
import numpy as np
import mediapipe as mp
from PIL import Image, ImageDraw, ImageFont

mpFaceMesh = mp.solutions.face_mesh
faceMesh = mpFaceMesh.FaceMesh(max_num_faces=1)

xiachunxia_ids = [400, 377, 152, 148, 176, 78, 95, 88, 178, 87, 14, 317, 402, 318, 324, 308, 400]

face_oval_ids = [10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288, 397, 365, 379, 378, 400, 377, 152, 148, 176,
                 149, 150, 136, 172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109, 10]

def cv2AddChineseText(img, text, position, textColor=(255, 0, 255), textSize=30):
    if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(img)
    # 字体的格式
    fontStyle = ImageFont.truetype(
        "SIMSUN.TTC", textSize, encoding="utf-8")
    # 绘制文本
    draw.text(position, text, textColor, font=fontStyle)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)


def get_faceLms(mp_FACE_MESH, frame):
    '''
    利用FaceMesh处理单帧图像，并返回关键点类
    :param frame: 单帧图像
    :return: faceLms类（包含各点坐标的集合）
    '''
    # frame.flags.writeable = False
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    results = mp_FACE_MESH.process(frame)
    # frame.flags.writeable = True
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    if results.multi_face_landmarks == None:
        return None
    faceLms = results.multi_face_landmarks[0]
    faceLms = results.multi_face_landmarks
    return faceLms


def cal_jianju(id1, id2, faceLms_landmark, h, w):
    zuo_zi = faceLms_landmark[id1]
    you_zi = faceLms_landmark[id2]
    x1, y1 = int(zuo_zi.x * w), int(zuo_zi.y * h)
    x2, y2 = int(you_zi.x * w), int(you_zi.y * h)
    return math.sqrt((x2-x1)**2+(y2-y1)**2)

def show_keypoint_pair(index1, index2, frame, faceLms, iw, ih, color):
    zijianju = cal_jianju(index1, index2, faceLms.landmark, ih, iw)
    x1, y1 = int(faceLms.landmark[index1].x * iw), int(faceLms.landmark[index1].y * ih)
    x2, y2 = int(faceLms.landmark[index2].x * iw), int(faceLms.landmark[index2].y * ih)
    cv2.circle(frame, (int(faceLms.landmark[index1].x * iw), int(faceLms.landmark[index1].y * ih)), 1, color=color,
               thickness=1)
    cv2.circle(frame, (int(faceLms.landmark[index2].x * iw), int(faceLms.landmark[index2].y * ih)), 1, color=color,
               thickness=1)
    cv2.line(frame, (x1, y1), (x2, y2), color=color, thickness=1)
    cv2.putText(frame, str(round(zijianju, 2)), (int(faceLms.landmark[index2].x * iw), int(faceLms.landmark[index2].y * ih)),
                cv2.FONT_HERSHEY_DUPLEX, 0.3, color=color, thickness=1)

    return frame

def show_keypoint_region(index_list, frame, faceLms, iw, ih):
    contours = []
    contour = []
    for id in index_list:
        x, y = int(faceLms.landmark[id].x * iw), int(faceLms.landmark[id].y * ih)
        contour.append((x, y))
    contours.append(contour)
    contours = np.array(contours)
    cv2.drawContours(frame, contours, -1, (0, 255, 0), 1)

    return frame

def show_keypoint_pairs(frame, iw, ih, mp_FACE_MESH):
    faceLms = get_faceLms(mp_FACE_MESH, frame)

    if faceLms == None:
        return frame

    for faceLm in faceLms:
        # 眼内眦间距
        frame = show_keypoint_pair(112, 362, frame, faceLm, iw, ih, color=(0, 0, 255))
        # 鼻翼宽度
        frame = show_keypoint_pair(64, 344, frame, faceLm, iw, ih, color=(0, 0, 255))

        # 口唇至下颌下缘
        frame = show_keypoint_pair(16, 152, frame, faceLm, iw, ih, color=(255, 0, 0))
        # 整个面部长度
        frame = show_keypoint_pair(9, 152, frame, faceLm, iw, ih, color=(255, 0, 0))

        # 上唇上缘至鼻小柱
        frame = show_keypoint_pair(2, 0, frame, faceLm, iw, ih, color=(0, 255, 0))
        # 下唇上缘至鼻小柱
        frame = show_keypoint_pair(2, 13, frame, faceLm, iw, ih, color=(0, 255, 0))

        # 下唇下面积
        frame = show_keypoint_region(xiachunxia_ids, frame, faceLm, iw, ih)
        # 脸部全面积
        frame = show_keypoint_region(face_oval_ids, frame, faceLm, iw, ih)

        # 外耳道口至颧骨前缘
        frame = show_keypoint_pair(93, 50, frame, faceLm, iw, ih, color=(0, 0, 255))
        # 外耳道口至下颌骨前缘
        frame = show_keypoint_pair(93, 152, frame, faceLm, iw, ih, color=(0, 0, 255))

    return frame

def show_keypoint_pairs_five(frame, iw, ih, mp_FACE_MESH):
    raw_frame = frame
    faceLms = get_faceLms(mp_FACE_MESH, raw_frame)

    if faceLms == None:
        return frame

    for faceLm in faceLms:
        # 眼内眦间距
        frame1 = show_keypoint_pair(112, 362, raw_frame.copy(), faceLm, iw, ih, color=(0, 0, 255))
        # 鼻翼宽度
        frame1 = show_keypoint_pair(64, 344, frame1, faceLm, iw, ih, color=(0, 0, 255))

        # 口唇至下颌下缘
        frame2 = show_keypoint_pair(16, 152, raw_frame.copy(), faceLm, iw, ih, color=(255, 0, 0))
        # 整个面部长度
        frame2 = show_keypoint_pair(9, 152, frame2, faceLm, iw, ih, color=(255, 0, 0))

        # 上唇上缘至鼻小柱
        frame3 = show_keypoint_pair(2, 0, raw_frame.copy(), faceLm, iw, ih, color=(0, 255, 0))
        # 下唇上缘至鼻小柱
        frame3 = show_keypoint_pair(2, 13, frame3, faceLm, iw, ih, color=(0, 255, 0))

        # 下唇下面积
        frame4 = show_keypoint_region(xiachunxia_ids, raw_frame.copy(), faceLm, iw, ih)
        # 脸部全面积
        frame4 = show_keypoint_region(face_oval_ids, frame4, faceLm, iw, ih)

        # 外耳道口至颧骨前缘
        frame5 = show_keypoint_pair(93, 50, raw_frame.copy(), faceLm, iw, ih, color=(0, 0, 255))
        # 外耳道口至下颌骨前缘
        frame5 = show_keypoint_pair(93, 152, frame5, faceLm, iw, ih, color=(0, 0, 255))

    return frame1, frame2, frame3, frame4, frame5
